"""处理训练和开发、测试数据集"""
from torch.utils.data import Dataset,DataLoader
from config import *
import pickle
import numpy as np

class TrainDataSet(Dataset):
    """
    训练数据集（大词林）
    """
    def __init__(self):
        with open(Config.word2vec_path, 'rb') as f:
            self.word2vec = pickle.load(f)
        self.wordpairs = []
        with open(Config.train_wordpairs_path,'r',encoding='utf-8') as f:
            for line in f:
                self.wordpairs.append(line.rstrip('\n').split())

    def __len__(self):
        return len(self.wordpairs)

    def __getitem__(self, idx):
        #return {'x':self.word2vec.get_word2vec(self.wordpairs[idx][0]),'y':np.concatenate((self.word2vec.get_word2vec(self.wordpairs[idx][1]),self.word2vec.get_word2vec(self.wordpairs[idx][2])))}
        return {'x': self.word2vec.word2id[self.wordpairs[idx][0]], 'y_1':self.word2vec.word2id[self.wordpairs[idx][1]], 'y_2':self.word2vec.word2id[self.wordpairs[idx][2]]}

class TestDataSet(Dataset):
    """
    测试数据集
    """
    def __init__(self):
        with open(Config.word2vec_path, 'rb') as f:
            self.word2vec = pickle.load(f)
        self.wordpairs = []
        with open(Config.test_wordpairs_path,'r',encoding='utf-8') as f:
            for line in f:
                self.wordpairs.append(line.rstrip('\n').split())

    def __len__(self):
        return len(self.wordpairs)

    def __getitem__(self, idx):
        return {'x': self.word2vec.word2id[self.wordpairs[idx][0]], 'y':self.word2vec.word2id[self.wordpairs[idx][1]], 'label':0 if self.wordpairs[idx][2]=="False" else 1}
